# coding=utf8

import matplotlib.pyplot as plt
import numpy as np
import tflearn
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.core import input_data, dropout, fully_connected
from tflearn.layers.estimator import regression

#train_x, train_y, test_x, test_y = tflearn.datasets.mnist.load_data(data_dir='/Users/vista/PycharmProjects/data/MNIST_data',one_hot=True)
dataset = tflearn.datasets.mnist.read_data_sets(train_dir='/Users/vista/PycharmProjects/data/MNIST_data',one_hot=True)

# train_x = train_x.reshape(-1, 28, 28, 1)
# test_x = test_x.reshape(-1, 28, 28, 1)

# 定义神经网络模型
conv_net = input_data(shape=[None, 28, 28, 1], name='input')
conv_net = conv_2d(conv_net, 32, 2, activation='relu')
conv_net = max_pool_2d(conv_net, 2)
conv_net = conv_2d(conv_net, 64, 2, activation='relu')
conv_net = max_pool_2d(conv_net, 2)
conv_net = fully_connected(conv_net, 1024, activation='relu')
conv_net = dropout(conv_net, 0.8)
conv_net = fully_connected(conv_net, 10, activation='softmax')
conv_net = regression(conv_net, optimizer='adam', loss='categorical_crossentropy', name='output')

model = tflearn.DNN(conv_net)

# 训练
# model.fit({'input': train_x}, {'output': train_y}, n_epoch=13,
#           validation_set=({'input': test_x}, {'output': test_y}),
#           snapshot_step=300, show_metric=True, run_id='mnist')

#model.save('/Users/vista/PycharmProjects/data/model/mnist/mnist.model')  # 保存模型

model.load('/Users/vista/PycharmProjects/data/model/mnist/mnist.model')   # 加载模型

fig = plt.figure()

for i in range(9):

    random_index = (np.random.rand(1) * dataset.test.num_examples)[0].astype(np.int32)

    x = dataset.test.images[random_index]
    xx = x.reshape([28,28])

    img = fig.add_subplot(3,3,i+1)
    img.imshow(xx, cmap='gray_r')
    img.set_xticks([])
    img.set_yticks([])

    a = model.predict(x.reshape(-1, 28, 28, 1))

    label = np.argmax( a , 1 )  # 预测

    plt.text(13, -1.2, str(label))

plt.show()

# ax4 = fig.add_subplot(212)
# ax4.imshow(x, cmap='gray_r')

#print model.evaluate(test_x,test_y)